# -*- coding: utf-8 -*-
# @Author: Tim Liu
# @Date: 2025-01-08
# @Last Modified by: Cursor
# @Last Modified time: 2025-01-08
import logging
# @Description: Neo4j Graph Engine Implementation

from typing import List, Dict, Any
from langchain_community.graphs.graph_document import GraphDocument
from langchain_neo4j import Neo4jGraph
from neo4j import GraphDatabase
from neo4j_graphrag.indexes import create_vector_index as neo4j_create_vector_index, upsert_vector as neo4j_upsert_vector
from config.settings import *
from .graph_engine import GraphEngine

class Neo4jEngine(GraphEngine):
    def __init__(self, url, user, password):
        """
        Initialize the Neo4jEngine with the given URL, user, and password. if not use default
        """
        default_url = NEO4J_GRAPH_ENDPOINT
        default_user = NEO4J_USER
        default_password = NEO4J_PASSWORD

        url = url or default_url
        user = user or default_user
        password = password or default_password
        logging.info("连接信息：url:{} user:{} password:{}".format(url, user, password))
        self.graph = Neo4jGraph(url, user, password, refresh_schema=False)
        logging.info("连接成功")
        self.driver = GraphDatabase.driver(url, auth=(user, password))

    def remove_existing_nodes(self, node_id: str = None, file_url: str = None):
        """
        Remove existing nodes by node_id or file_url.
        """
        if node_id:
            query = f"MATCH (n) WHERE n.id = '{node_id}' DETACH DELETE n"
        elif file_url:
            query = f"MATCH (n) WHERE n.source_url = '{file_url}' DETACH DELETE n"
        else:
            raise ValueError("Either node_id or file_url must be provided")
        
        self.graph.query(query)

    def add_graph_documents(self, graph_documents: List[GraphDocument]):
        """
        Add graph documents to the Neo4j database.
        """
        self.graph.add_graph_documents(graph_documents)

    def run_query(self, query: str) -> List[Dict[str, Any]]:
        """
        Run a query and return the results.
        """
        return self.graph.query(query)
    
    def get_nodes_by_type(self, node_type: str) -> List[Dict[str, Any]]:
        """
        Get nodes by type from the Neo4j database.

        Args:
            node_type (str): The type of nodes to retrieve.

        Returns:
            List[Dict[str, Any]]: A list of nodes of the specified type.
        """
        query = f"MATCH (n:{node_type}) RETURN n"
        #result = self.graph.query(query)
        with self.driver.session() as session:
            result = session.run(query)
            return [record['n'] for record in result]

    def create_vector_index(self, index_name: str, label: str, embedding_property: str, dimensions: int, similarity_fn: str, fail_if_exists: bool):
        """
        Create a vector index in the Neo4j database.
        """
        neo4j_create_vector_index(
            self.driver,
            index_name,
            label=label,
            embedding_property=embedding_property,
            dimensions=dimensions,
            similarity_fn=similarity_fn,
            fail_if_exists=fail_if_exists
        )

    def upsert_vector(self, node_id: str, embedding_property: str, vector: List[float]):
        """
        Upsert a vector for a node in the Neo4j database.
        """
        neo4j_upsert_vector(
            self.driver,
            node_id=node_id,
            embedding_property=embedding_property,
            vector=vector
        )